package org.ruoyi.knowledge.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Request;
import okhttp3.HttpUrl;
import org.ruoyi.common.core.utils.OkHttpUtil;
import org.ruoyi.knowledge.config.DifyProperties;
import org.ruoyi.knowledge.service.IDifyKnowledgeService;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class DifyKnowledgeServiceImpl implements IDifyKnowledgeService {

    private final OkHttpUtil okHttpUtil;
    private final DifyProperties difyProperties;
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public String retrieve(String datasetId, String query, Integer topK) {
        try {
            okHttpUtil.setApiHost(stripTrailingSlash(difyProperties.getApiHost()));
            okHttpUtil.setApiKey("Bearer " + difyProperties.getApiKey());

            Map<String, Object> payload = new HashMap<>();
            payload.put("query", query);
            if (topK != null) {
                payload.put("top_k", topK);
            }

            String json = objectMapper.writeValueAsString(payload);
            String url = "/v1/datasets/" + datasetId + "/retrieval";
            Request request = okHttpUtil.createPostRequest(url, json);
            return okHttpUtil.executeRequest(request);
        } catch (Exception e) {
            log.error("Dify retrieve error: {}", e.getMessage(), e);
            return null;
        }
    }

    @Override
    public String listDatasets(Integer page, Integer pageSize) {
        try {
            okHttpUtil.setApiHost(stripTrailingSlash(difyProperties.getApiHost()));
            okHttpUtil.setApiKey("Bearer " + difyProperties.getApiKey());

            HttpUrl.Builder builder = HttpUrl.parse(difyProperties.getApiHost() + "/v1/datasets").newBuilder();
            if (page != null) builder.addQueryParameter("page", String.valueOf(page));
            if (pageSize != null) builder.addQueryParameter("limit", String.valueOf(pageSize));

            Request request = new Request.Builder()
                .url(builder.build())
                .header("Authorization", "Bearer " + difyProperties.getApiKey())
                .build();
            return okHttpUtil.executeRequest(request);
        } catch (Exception e) {
            log.error("Dify listDatasets error: {}", e.getMessage(), e);
            return null;
        }
    }

    @Override
    public String listDocuments(String datasetId, Integer page, Integer pageSize) {
        try {
            okHttpUtil.setApiHost(stripTrailingSlash(difyProperties.getApiHost()));
            okHttpUtil.setApiKey("Bearer " + difyProperties.getApiKey());

            HttpUrl.Builder builder = HttpUrl.parse(difyProperties.getApiHost() + "/v1/datasets/" + datasetId + "/documents").newBuilder();
            if (page != null) builder.addQueryParameter("page", String.valueOf(page));
            if (pageSize != null) builder.addQueryParameter("limit", String.valueOf(pageSize));

            Request request = new Request.Builder()
                .url(builder.build())
                .header("Authorization", "Bearer " + difyProperties.getApiKey())
                .build();
            return okHttpUtil.executeRequest(request);
        } catch (Exception e) {
            log.error("Dify listDocuments error: {}", e.getMessage(), e);
            return null;
        }
    }

    private String stripTrailingSlash(String host) {
        if (host == null) return null;
        if (host.endsWith("/")) return host.substring(0, host.length() - 1);
        return host;
    }
}


